import glob
import math
import os.path

import numpy as np
from osgeo import gdal


# 写入tif
def write_tif(file_path, data, geotransform, nodata, dataType):
    gdal_type = ''
    if dataType == 'int16':
        gdal_type = gdal.GDT_Int16
    elif dataType == 'int32':
        gdal_type = gdal.GDT_Int32
    elif dataType == 'float32':
        gdal_type = gdal.GDT_Float32
    elif dataType == 'float64':
        gdal_type = gdal.GDT_Float64
    elif dataType == 'byte':
        gdal_type = gdal.GDT_Byte
    elif dataType == 'uint16':
        gdal_type = gdal.GDT_UInt16
    elif dataType == 'uint32':
        gdal_type = gdal.GDT_UInt32
    # 相当于一个创建数据集的驱动
    driver = gdal.GetDriverByName("GTiff")
    # 根据数据的维度来确定行列数（图像大小）
    rows, cols = data.shape
    # 创建一个新的数据集，存储输出文件；1是波段
    dataset = driver.Create(file_path, cols, rows, 1, gdal_type)
    # 设置仿射变换矩阵
    dataset.SetGeoTransform(geotransform)
    # 设置投影
    prj = osr.SpatialReference()
    # WGS84投影
    prj.ImportFromEPSG(4326)
    dataset.SetProjection(prj.ExportToWkt())
    # 1个波段
    band = dataset.GetRasterBand(1)
    band.WriteArray(data)
    band.SetNoDataValue(nodata)
    # 释放内存
    del dataset


def mean_filter(image, kernel_size=3):
    """
    均值滤波
    :param image:
    :param kernel_size:
    :return:
    """
    h, w = image.shape
    filterd = np.zeros((h, w))
    half_kernel_size = kernel_size // 2

    for i in range(half_kernel_size, h - half_kernel_size):
        for j in range(half_kernel_size, w - half_kernel_size):
            filterd[i, j] = np.nanmean(image[i - half_kernel_size:i + half_kernel_size + 1, j - half_kernel_size:j + half_kernel_size + 1])
    return filterd


# 计算太阳校正系数
def get_sun_correction_factor(phi, nd):
    # 公式： C = sinφsinδ(1 - tan²φtan²δ)^0.5 + cosφcosδarccos(-tanφtanδ)
    # δ = 0.006918 - 0.399912cos(Γ) + 0.070257sin(Γ) - 0.006758cos(2Γ) + 0.000907sin(2Γ) - 0.002697cos(3Γ) + 0.00148sin(3Γ)
    # Γ = (2 * π * (nd - 1)) / 365.25
    # φ是纬度 δ为太阳赤纬 Γ为昼角 nd为日数
    gamma = (2 * np.pi * (nd - 1)) / 365.25
    delta = 0.006918 - 0.399912 * np.cos(gamma) + 0.070257 * np.sin(gamma) \
            - 0.006758 * np.cos(2 * gamma) + 0.000907 * np.sin(2 * gamma) \
            - 0.002697 * np.cos(3 * gamma) + 0.00148 * np.sin(3 * gamma)
    # 百度百科
    # delta = np.arcsin(0.39795 * np.cos(0.98563 * (nd - 173)))
    print(gamma)
    print(delta * 180 / np.pi)
    C = np.sin(phi) * np.sin(delta) * (1 - np.tan(phi) ** 2 * np.tan(delta) ** 2) ** 0.5 + \
        np.cos(phi) * np.cos(delta) * np.arccos(-np.tan(phi) * np.tan(delta))
    return C


# 计算ati
def get_ati(day_lst, night_lst, albedo, doy, geo):
    # 对数据进行筛选，剔除不满足计算要求的数据，主要有：
    # nodata数据、异常值数据
    day_lst[day_lst <= 0] = np.nan
    night_lst[night_lst <= 0] = np.nan
    albedo[albedo <= 0] = np.nan
    albedo[albedo == -9999] = np.nan

    # 实现太阳校正系数校正
    lat_arr = np.zeros_like(day_lst, dtype=np.float32)
    for i in range(day_lst.shape[0]):
        for j in range(day_lst.shape[1]):
            lat_arr[i, j] = math.radians(geo[3] + geo[4] * i + geo[5] * j)

    C = get_sun_correction_factor(lat_arr, doy)

    # 计算ati
    ati = C * (1 - albedo * 0.0001) / (day_lst - night_lst)
    # WMA滤波
    ati = mean_filter(ati, 3)
    # 剔除异常值
    ati[ati > 1] = np.nan
    ati[ati < 0] = np.nan
    return ati


if __name__ == '__main__':
    day_lst_dir = r"白天LST数据目录"
    night_lst_dir = r"夜晚LST数据目录"
    albedo_dir = r"albedo数据目录"
    # 获取数据路径列表
    day_lst_path_lst = glob.glob(day_lst_dir + "/*.tif")
    night_lst_path_lst = glob.glob(night_lst_dir + "/*.tif")
    albedo_path_lst = glob.glob(albedo_dir + "/*.tif")
    output_dir = r"输出目录"
    for day_lst_path, night_lst_path, albedo_path in zip(day_lst_path_lst, night_lst_path_lst, albedo_path_lst):
        filename = os.path.basename(day_lst_path)
        doy = int(filename.split(".")[0][-3:])
        day_lst_dataset = gdal.Open(day_lst_path)
        night_lst_dataset = gdal.Open(night_lst_path)
        albedo_dataset = gdal.Open(albedo_path)
        geo = day_lst_dataset.GetGeoTransform()

        day_lst = day_lst_dataset.GetRasterBand(1).ReadAsArray()
        night_lst = night_lst_dataset.GetRasterBand(1).ReadAsArray()
        albedo = albedo_dataset.GetRasterBand(1).ReadAsArray()

        ati = get_ati(day_lst, night_lst, albedo, doy, geo)
        output_path = os.path.join(output_dir, filename)
        write_tif(output_path, ati, geo, nodata=-9999, dataType='float32')
